from typing import Any, Callable, TYPE_CHECKING, Union

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, resume
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .source import parquet_file_source_context

if TYPE_CHECKING:
    from pyarrow.parquet import ParquetFile, RowGroupMetaData


class ParquetFileRowGroupDataset(Dataset[dict]):
    def __init__(self,
        source: Union[Callable[[], 'ParquetFile'], 'ParquetFile'],
        row_group_index: int,
        metadata: 'RowGroupMetaData',
    ):
        self._source = source
        self._row_group_index = row_group_index
        self._metadata = metadata

    @property
    def metadata(self) -> 'RowGroupMetaData':
        return self._metadata

    def _load_samples(self):
        with parquet_file_source_context(self._source) as parquet_file:
            row_group = parquet_file.read_row_group(self._row_group_index)
        return wrap_dataset(row_group.to_pylist())

    def __len__(self) -> int:
        return self.metadata.num_rows

    def __getitem__(self, index: int, /) -> dict:
        return self._load_samples()[index]

    def __iter__(self) -> DatasetIterator[dict]:
        return iter(self._load_samples())

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[dict]:
        return resume(self._load_samples(), state)
